import numpy as np
from pyE17 import utils as U
from pyE17 import io
import pyE17
import ptycho
import os
import matplotlib.cm as cm
import matplotlib.pyplot as plt

# Main parameters
energy = 6.2
z = 7.046        # Distance chosen to have exactly 32 nm pixel size
ds = 172e-6
lam = 1.2398e-9 / energy

# Number of incident photons in the probe
noise = 3e8

# Carbon, thick
ior = -1j * 3.4054267E-08 - 1.19337546E-05
max_thickness = 1e-6

run_name = 'lownoise_sp_displacements'

subpix_method = 'linear'

# Displacement rms in pixels
disp_rms = 1. 

# subpix functions are not exposed in ptycho (maybe this should be done)
def subpix_fourier(a, ramp, forward=True):
    """\
    Perform a shift with the provided ramp.
    """
    if forward:
        return np.fft.ifftn(np.fft.fftn(a) * ramp)
    else:
        return np.fft.ifftn(np.fft.fftn(a) * ramp.conj())

def subpix_linear(a, pts, forward=True):
    """\
    Perform a subpixel shift using bilinear interpolation.
    """
    x0, x1 = pts
    if not forward: x0, x1 = -x0, -x1
    a00 = U.pshift(a, [np.floor(x0), np.floor(x1)], 'nearest')
    a01 = U.pshift(a, [np.floor(x0), np.ceil(x1)], 'nearest')
    a10 = U.pshift(a, [np.ceil(x0), np.floor(x1)], 'nearest')
    a11 = U.pshift(a, [np.ceil(x0), np.ceil(x1)], 'nearest')
    x0 = x0 % 1
    x1 = x1 % 1
    return a00 * (1 - x0) * (1 - x1) + a01 * (1 - x0) * x1 + a10 * x0 * (1 - x1) + a11 * x0 * x1


probe_file = 'probe_0.h5'
object_file = 'object_0.h5'

probe = io.h5read(probe_file)['probe']
probe /= U.norm(probe)
ob0 = io.h5read(object_file)['object']
object = np.exp(-2j * np.pi * ior * max_thickness * (255 - ob0) / 255. / lam)

asize = probe.shape
dx_spec = lam * z / (asize[0] * ds) 

nth = 5
dr = 0.7e-6
lx = 12e-6
ly = 12e-6

positions = ptycho.round_scan_ROI_positions(dr, lx, ly, nth)

positions = np.array(positions) / dx_spec
positions -= positions.min(axis=0)
positions_float = positions
positions = np.round(positions_float)
Npos = len(positions)

# Create the random displacements
displacements = disp_rms * np.random.normal(size=positions.shape)

# Here we can combine the displacements in subpix_parameters
if subpix_method.lower() == 'fourier':
    p_shift = positions - positions_float - displacements
    subpix_parameter = [np.exp(2j * np.pi * np.sum(U.fgrid(asize, x / asize), axis=0)) for x in p_shift]
elif subpix_method.lower() == 'linear':
    subpix_parameter = positions - positions_float - displacements

# This is the way the final object size is computed 
object_size = tuple(np.array(asize) + positions.max(axis=0))

# Setup object views for computation of diffraction patterns
offset = (300, 300)
obj = object[offset[0]:(offset[0] + object_size[0]), offset[1]:(offset[1] + object_size[1])]
obview = [obj[p[0]:(p[0] + asize[0]), p[1]:(p[1] + asize[1])] for p in positions]
coverage = np.zeros_like(obj).astype(float)
for p in positions:
    coverage[p[0]:(p[0] + asize[0]), p[1]:(p[1] + asize[1])] += abs(probe) ** 2

io.h5write('simul_%s_solution.h5' % run_name, obj=obj.copy(), probe=probe, coverage=coverage)

Iprobe = (np.abs(np.fft.fftn(probe)) ** 2).sum()

data = np.zeros((Npos,) + asize, dtype=int)
for ii, iob in enumerate(range(len(obview))):
    if subpix_method == 'fourier':
        pr = subpix_fourier(probe, subpix_parameter[ii])
    elif subpix_method == 'linear':
        pr = subpix_linear(probe, subpix_parameter[ii])
        
    probe_big = np.zeros_like(obj)
    posi = positions[ii]
    probe_big[posi[0]:(posi[0] + asize[0]), posi[1]:(posi[1] + asize[1])] = pr
    fac = probe_big * obj
    p = U.imsave(fac, 'probe_object_%d.png' % ii)
    plt.imsave('probe_object_%d_bw.png' % ii, fac, cmap=cm.binary)
    dp = np.fft.fftshift(np.fft.fftn(pr * obview[iob]))
    p = U.imsave(abs(np.power(dp, 0.25)), 'probe_object_fft_%d.png' % ii)
    data[iob, :, :] = np.random.poisson(np.abs(dp) ** 2 * noise / Iprobe)
    p = U.imsave(np.power(data[iob, :, :], 0.25), 'possion_noise_%d.png' % ii)
    print ii
    
dirname = 'simul_%s' % run_name
try:
    os.mkdir(dirname)
except OSError:
    pass
filename = dirname + ('/simul_%s_data_%dx%d.h5' % (run_name, asize[0], asize[1]))
print filename
io.h5write(filename, data=data, fmask=np.ones(asize, dtype=int), displacements=displacements)
